import cv2
import numpy as np

class PVD:
    def __init__(self):
        # 固定范围表，确保嵌入提取一致
        self.range_table = [
            (0, 7, 3),      # 0-7: 3 bits
            (8, 15, 3),     # 8-15: 3 bits  
            (16, 31, 4),    # 16-31: 4 bits
            (32, 63, 5),    # 32-63: 5 bits
            (64, 127, 6),   # 64-127: 6 bits
            (128, 255, 7)   # 128-255: 7 bits
        ]
        
    def get_range_info(self, difference):
        """获取差值对应的范围信息"""
        for lower, upper, bits in self.range_table:
            if lower <= difference <= upper:
                return lower, bits
        return self.range_table[0][0], self.range_table[0][2]  # 默认第一个范围

    def embed(self, image_path, message, output_path, channel=0):
        try:
            print("=== PVD嵌入开始 ===")
            print(f"消息: {message}")
            print(f"消息长度: {len(message)} 字符")
            
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print("错误：无法读取图像")
                return False
            
            print(f"图像尺寸: {img.shape}")
            print(f"使用通道: {channel}")
            
            # 添加明确的结束标记
            message_with_end = message + "END"  # 简单的结束标记
            print(f"完整消息: {message_with_end}")
            
            # 转换为二进制（使用UTF-8编码）
            binary_data = ""
            for char in message_with_end:
                binary_data += format(ord(char), '08b')  # 每个字符8位
            
            print(f"二进制数据: {binary_data}")
            print(f"二进制长度: {len(binary_data)} 位")
            
            # 嵌入数据
            height, width = img.shape[:2]
            embedded_img = img.copy()
            bit_index = 0
            total_bits = len(binary_data)
            
            print(f"开始嵌入，总需要 {total_bits} 位")
            
            for i in range(height):
                for j in range(0, width - 1, 2):  # 水平像素对
                    if bit_index >= total_bits:
                        break
                    
                    # 获取当前像素对
                    p1 = int(embedded_img[i, j, channel])
                    p2 = int(embedded_img[i, j+1, channel])
                    
                    # 计算原始差值
                    original_diff = abs(p1 - p2)
                    lower_bound, bits_capacity = self.get_range_info(original_diff)
                    
                    # 确定要嵌入的位数
                    bits_to_embed = min(bits_capacity, total_bits - bit_index)
                    if bits_to_embed <= 0:
                        continue
                    
                    # 提取要嵌入的位
                    secret_bits = binary_data[bit_index:bit_index + bits_to_embed]
                    secret_value = int(secret_bits, 2)
                    
                    # 计算新的差值
                    new_diff = lower_bound + secret_value
                    
                    # 调整像素值
                    if p1 >= p2:
                        change = new_diff - original_diff
                        new_p1 = p1 + (change + 1) // 2
                        new_p2 = p2 - change // 2
                    else:
                        change = new_diff - original_diff
                        new_p1 = p1 - change // 2
                        new_p2 = p2 + (change + 1) // 2
                    
                    # 边界检查
                    new_p1 = max(0, min(255, new_p1))
                    new_p2 = max(0, min(255, new_p2))
                    
                    # 更新图像
                    embedded_img[i, j, channel] = new_p1
                    embedded_img[i, j+1, channel] = new_p2
                    
                    bit_index += bits_to_embed
                    
                    if bit_index % 100 == 0:
                        print(f"已嵌入 {bit_index}/{total_bits} 位")
            
            print(f"嵌入完成! 总共嵌入 {bit_index} 位")
            
            # 保存为PNG格式
            if not output_path.lower().endswith('.png'):
                output_path = output_path.rsplit('.', 1)[0] + '.png'
                
            success = cv2.imwrite(output_path, embedded_img)
            if success:
                print(f"图像保存成功: {output_path}")
                return True
            else:
                print("图像保存失败")
                return False
                
        except Exception as e:
            print(f"PVD嵌入错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return False

    def extract(self, image_path, channel=0):
        try:
            print("=== PVD提取开始 ===")
            
            # 读取图像
            img = cv2.imread(image_path)
            if img is None:
                print("错误：无法读取图像")
                return "无法读取图像"
            
            print(f"图像尺寸: {img.shape}")
            print(f"使用通道: {channel}")
            
            # 提取二进制数据
            binary_data = ""
            height, width = img.shape[:2]
            
            print("开始提取二进制数据...")
            
            for i in range(height):
                for j in range(0, width - 1, 2):  # 水平像素对
                    # 获取当前像素对
                    p1 = int(img[i, j, channel])
                    p2 = int(img[i, j+1, channel])
                    
                    # 计算当前差值
                    current_diff = abs(p1 - p2)
                    lower_bound, bits_capacity = self.get_range_info(current_diff)
                    
                    # 提取秘密值
                    secret_value = current_diff - lower_bound
                    if secret_value < 0:
                        continue
                    
                    # 转换为二进制
                    secret_bits = format(secret_value, f'0{bits_capacity}b')
                    binary_data += secret_bits
            
            print(f"提取的二进制数据长度: {len(binary_data)} 位")
            print(f"二进制数据: {binary_data}")
            
            # 转换为字符串
            result_chars = []
            temp_byte = ""
            
            for bit in binary_data:
                temp_byte += bit
                if len(temp_byte) == 8:
                    try:
                        char_code = int(temp_byte, 2)
                        char = chr(char_code)
                        result_chars.append(char)
                        temp_byte = ""
                        
                        # 检查结束标记
                        current_text = ''.join(result_chars)
                        if current_text.endswith("END"):
                            final_result = current_text[:-3]  # 移除"END"
                            print(f"找到结束标记! 最终结果: {final_result}")
                            return final_result
                    except:
                        temp_byte = ""
                        continue
            
            # 如果没有找到结束标记
            final_result = ''.join(result_chars)
            print(f"未找到结束标记，提取的数据: {final_result}")
            return final_result if final_result else "未提取到有效信息"
            
        except Exception as e:
            print(f"PVD提取错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return f"提取错误: {str(e)}"